/****************************************************************************
 *
 * Copyright (C) 2003-2008 Los Alamos National Security, LLC
 *                         Packet Analytics Corporation
 *
 * This program is free software; you can redistribute it and/or modify
 * it under the terms of the GNU General Public License Version 2 as
 * published by the Free Software Foundation.  You may not use, modify or
 * distribute this program under any other version of the GNU General
 * Public License.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program; if not, write to the Free Software
 * Foundation, Inc., 59 Temple Place - Suite 330, Boston, MA 02111-1307, USA.
 *
 ****************************************************************************/
package nfse.dataset.proofpoint;

import java.sql.ResultSet;
import java.sql.Statement;
import java.sql.Timestamp;
import java.util.*;
import java.io.*;

import nfse.JDBCDatabase;
import nfse.NetFSEUtil;

import java.lang.Thread;
import java.net.Socket;

import javax.net.SocketFactory;
import javax.net.ssl.SSLSocketFactory;

public class ProofPointBuilder {

	private static String server = "";

	private static int serverPort = 9097;

	private static boolean useSSL = false;

	private static String ppDirPath = "";

	private static int sensor = 7;

	private static String startDate = "2005-01-01";

	private static String endDate = "";

	private static String keyword = "";

	private static String buildListPath = "";

	private static String dbDatabase = "proofpoint";

	private static String dbServer = "localhost";

	private static String dbUser = "root";

	private static String dbPassword = "metadata";
	private static String dbPort = "3306";

	private static boolean stdin = false;

	public static void main(String[] args) {
		try {
			if (args.length != 1) {
				System.out
						.println("Usage: java ProofPointBuilder config_file_path");
				return;
			}

			File conf = new File(args[0]);
			if (!conf.exists()) {
				System.err.println("Could not find the specified config file: "
						+ args[0]);
				System.exit(1);
			}

			LineNumberReader in = new LineNumberReader(new FileReader(conf));
			String line = in.readLine();
			while (line != null) {
				line = line.trim();
				if (line.length() > 0) {
					String[] tokens = line.split(" +");
					if (tokens[0].equals("server")) {
						server = tokens[1];
					} else if (tokens[0].equals("build_list_path")) {
						buildListPath = tokens[1];
					} else if (tokens[0].equals("log_dir")) {
						ppDirPath = tokens[1];
					} else if (tokens[0].equals("start_date")) {
						startDate = tokens[1];
					} else if (tokens[0].equals("end_date")) {
						endDate = tokens[1];
					} else if (tokens[0].equals("log_keyword")) {
						keyword = tokens[1];
					} else if (tokens[0].equals("sensor")) {
						sensor = Integer.parseInt(tokens[1]);
					} else if (tokens[0].equals("server_port")) {
						serverPort = Integer.parseInt(tokens[1]);
					} else if (tokens[0].equals("use_stdin")) {
						stdin = true;
					} else if (tokens[0].equals("database")) {
						dbServer = tokens[1];
						dbDatabase = tokens[2];
						dbUser = tokens[3];
						dbPassword = tokens[4];
					}
				}
				line = in.readLine();
				// System.out.println(line);
			}
			in.close();

			if (stdin) {
				BufferedReader reader = new BufferedReader(
						new InputStreamReader(System.in));
				insertProofPoint(reader);
				reader.close();
			} else {

				// String tempFilePath = args[2];
				while (true) {
					Vector<ProofPointLog> temp = new Vector<ProofPointLog>();
					in = new LineNumberReader(new FileReader(buildListPath));
					line = in.readLine();
					while (line != null) {
						ProofPointLog log = ProofPointLog.generate(line);
						if (log != null) {
							temp.addElement(log);
						}
						line = in.readLine();
					}
					in.close();

					ProofPointLog[] entries = new ProofPointLog[temp.size()];
					for (int i = 0; i < entries.length; i++) {
						entries[i] = (ProofPointLog) temp.elementAt(i);
					}
					// Arrays.sort(entries, new LFAPLogComparator());
					Arrays.sort(entries);

					// LFAPFileFilter filter = new LFAPFileFilter();
					File ppDir = new File(ppDirPath);
					if (!ppDir.isDirectory()) {
						System.err.println("Invalid directory specified: "
								+ ppDirPath);
						System.exit(1);
					}
					File[] files = ppDir.listFiles(new ProofPointFileFilter(
							startDate, endDate));
					if (files == null) {
						System.err
								.println("No ProofPoint files! This is rather bad");
						System.exit(0);
					}
					// Arrays.sort(files, new FileComparator());
					Arrays.sort(files);
					// System.out.println("Num files=" + files.length);
					boolean didSomething = false;
					for (int i = files.length - 1; i >= 0; i--) {
						if (didSomething) {
							break;
						}
						String filePath = files[i].getAbsolutePath();
						System.out.println(filePath);
						boolean found = false;
						for (int j = 0; j < entries.length; j++) {
							if (filePath.equals(entries[j].filePath)) {
								found = true;
								break;
							}
						}
						if (!found) { // Need to process
							String exec = "";
							int val = 0;
							didSomething = true;
							ProofPointLog log = ProofPointLog
									.generate(files[i]);
							String logDateStr = log.getDateStr();
							boolean withinDateRange = true;

							if (withinDateRange) {

								File f = new File(log.filePath);
								System.out.println(f.getAbsolutePath());
								String fileName = f.getName();
								// System.out.println(fileName);
								long size = 0;
								do {
									if (size != 0)
										System.out
												.println("Waiting while being copied.");
									size = f.length();
									Thread.sleep(5000);
									// Process p =
									// Runtime.getRuntime().exec("sleep
									// 5");
									// p.waitFor();
								} while (f.length() != size); // Checks that
								// the
								// file is not being
								// copied.

								BufferedReader reader = new BufferedReader(
										new InputStreamReader(
												new FileInputStream(f)));

								insertProofPoint(reader);
								reader.close();
							} else {
								System.out.println("Skipping: "
										+ log.toString());
							}
							if (val != 0) {
								System.out.println("Failed command line: "
										+ exec);
								Process p2 = Runtime.getRuntime().exec(
										"sleep 60");
								int val2 = p2.waitFor();
							} else {
								PrintWriter out = new PrintWriter(
										new FileWriter(args[0]));
								for (int x = 0; x < entries.length; x++)
									out.println(entries[x]);
								out.println(log);
								out.close();
							}
						}
					}
					if (!didSomething) {
						System.out.println("Sleeping...");
						Process p = Runtime.getRuntime().exec("sleep 30");
						p.waitFor();
					}
				}
			}
		} catch (Exception E) {
			E.printStackTrace();
		}
	}

	private static void insertProofPoint(BufferedReader in) {
		try {
			SocketFactory sf = null;
			if (useSSL)
				sf = SSLSocketFactory.getDefault(); // Use SSL sockets
			else
				sf = SocketFactory.getDefault(); // Use a plaintext socket
			Socket sock = sf.createSocket(server, serverPort);

			DataInputStream dis = new DataInputStream(sock.getInputStream());
			DataOutputStream dos = new DataOutputStream(sock.getOutputStream());
			// BufferedReader br = new BufferedReader(new
			// InputStreamReader(sock.getInputStream()));

			dos.writeInt(sensor);
			dos.writeInt(1);
			dos.writeInt(1);
			dos.writeInt(1);
			dos.writeInt(1);
			dos.writeInt(7); // month
			dos.writeInt(7); // day
			dos.writeInt(2007); // year
			dos.writeInt(1); // run

			GregorianCalendar cal = new GregorianCalendar();
			cal.add(Calendar.MONTH, -1);

			System.out.println("Connecting to jdbc:mysql://" + dbServer + ":"
					+ dbPort + "/" + dbDatabase + ", " + dbUser + ", "
					+ dbPassword);

			// Connect to the remote ProofPoint database
			JDBCDatabase db = new JDBCDatabase("jdbc:mysql://" + dbServer + ":"
					+ dbPort + "/" + dbDatabase, dbUser, dbPassword);
			Statement s = db.createStatement();
			Statement s2 = db.createStatement();

			s.execute("delete from proofpoint_temp where TS < '"
					+ new Timestamp(cal.getTimeInMillis()) + "'");

			// LineNumberReader in = new LineNumberReader(new FileReader(file));
			String line = in.readLine();
			while (line != null) {

				if (keyword.length() == 0
						|| (line.indexOf(keyword) > 0 && line.indexOf("cmd=") > 0)) {
					// System.out.println(line);
					try {
						line = line.substring(line.indexOf(" s=") + 1);
						line = line.trim();
						String[] tokens = line.split(" ");
						Timestamp ts = new Timestamp(System.currentTimeMillis());
						if (tokens[tokens.length - 1].startsWith("start=")) {
							int idx1 = tokens.length - 1;
							int idx2 = tokens[idx1].indexOf("=");
							ts = new Timestamp(Long.parseLong(tokens[idx1]
									.substring(idx2 + 1)) * 1000);
						}
						String sessionID = tokens[0].substring(2);
						String cmd = "";
						if (tokens[1].startsWith("m="))
							cmd = tokens[3];
						else
							cmd = tokens[2];

						if (cmd.equals("cmd=disconnect")) {
							// System.out.println(line);
							int avIdx = -1;
							int spamIdx = -1;
							Vector<String> records = new Vector<String>();
							int count = 0;
							String ip = null;
							String country = "";
							String virusName = "";
							String subject = " ";
							Vector<String> recipients = new Vector<String>();
							String sender = "";
							Timestamp startTS = null;
							StringBuffer recordBlob = new StringBuffer();

							String sql = "select min(TS) a from proofpoint_temp where Session_ID='"
									+ sessionID + "'";
							ResultSet rs = s2.executeQuery(sql);
							if (rs.next()) {
								startTS = rs.getTimestamp("a");
							}
							rs.close();

							sql = "select * from proofpoint_temp where Session_ID='"
									+ sessionID + "'";
							rs = s2.executeQuery(sql);
							while (rs.next()) {
								String recordHex = rs.getString("Record");
								String record = NetFSEUtil
										.convertFromHex(recordHex);
								records.addElement(record);
								recordBlob.append(recordHex + " ");
								// if ((startTS == null) &&
								// (record.indexOf("cmd=connect") > 0)) {
								// startTS = rs.getTimestamp("TS");
								// }
								if ((record.indexOf(" module=av ") > 0)
										&& (record.indexOf(" rule=notcleaned ") > 0)) {
									avIdx = count; // We found the indicator
									// that
									// there
									// was a virus in this email
								} else if (record.indexOf("rule=spam") > 0) {
									spamIdx = count;
									int subjIdx = record.indexOf("subject=\"");
									if (subjIdx != -1) {
										int qIdx1 = record.indexOf("\"",
												subjIdx);
										if (qIdx1 != -1) {
											int qIdx2 = record.indexOf("\"",
													qIdx1 + 1);
											subject = record.substring(
													qIdx1 + 1, qIdx2);
										}
									}
								}
								count++;
							}
							rs.close();

							// If this is an email with a virus, process it
							if (avIdx > -1 || spamIdx > -1) {
								// System.out.println(line);

								for (int i = 0; i < records.size(); i++) {
									String record = (String) records
											.elementAt(i);
									// System.out.println(record);
									if (record.indexOf("cmd=connect") > 0) {
										// This is the connect record
										// Find the IP and country if possible
										int ipIdx = record.indexOf("ip=");
										ip = record.substring(ipIdx + 3, record
												.indexOf(" ", ipIdx));
										int countryIdx = record
												.indexOf("country=");
										if (countryIdx > 0)
											country = record
													.substring(countryIdx + 8,
															record.indexOf(" ",
																	countryIdx));
									} else if ((record.indexOf("mod=av") > 0)
											&& (record
													.indexOf("rule=notcleaned") > 0)) {
										// This is the record containing the
										// virus
										// name
										// System.out.println(record);
										int nameIdx = record.indexOf("name=");
										int nextSpace = record.indexOf(" ",
												nameIdx);
										if (nextSpace > 0)
											virusName = record.substring(
													nameIdx + 5, nextSpace);
										else
											virusName = record
													.substring(nameIdx + 5);
									} else if (record.indexOf("cmd=env_rcpt") > 0) {
										// This indicates one of the recipients
										// of
										// the
										// email
										int valueIdx = record.indexOf("value=");
										int nextSpace = record.indexOf(" ",
												valueIdx);
										if (nextSpace > 0)
											recipients.addElement(record
													.substring(valueIdx + 6,
															nextSpace));
										else
											recipients.addElement(record
													.substring(valueIdx + 6));
									} else if (record.indexOf("cmd=env_from") > 0) {
										// This record contains the sender
										int valueIdx = record.indexOf("value=");
										int nextSpace = record.indexOf(" ",
												valueIdx);
										if (nextSpace > 0)
											sender = record.substring(
													valueIdx + 6, nextSpace);
										else
											sender = record
													.substring(valueIdx + 6);
									}

								}
								recordBlob
										.append(NetFSEUtil.convertToHex(line));

								// System.out.println("VirusName='" + virusName
								// + "'");
								if (ip != null && !ip.equals("null")) {
									virusName = virusName.replaceAll(",", " ");
									subject = subject.replaceAll(",", " ");
									virusName = virusName.replaceAll("\\", " ");
									subject = subject.replaceAll("\\", " ");
									for (int i = 0; i < recipients.size(); i++) {
										String logStr = ""
												+ sessionID
												+ ","
												+ startTS
												+ ","
												+ ip
												+ ","
												+ sender
												+ ","
												+ (String) recipients
														.elementAt(i) + ","
												+ virusName + "," + subject;
										// System.out.println(logStr);
										byte[] bytes = logStr.getBytes();
										dos.writeInt(bytes.length);
										dos.write(bytes);
									}
								}
							}

							s2
									.execute("delete from proofpoint_temp where Session_ID='"
											+ sessionID + "'");

						} else {
							String sql = "insert into proofpoint_temp values ('"
									+ sessionID
									+ "','"
									+ ts
									+ "','"
									+ NetFSEUtil.convertToHex(line) + "')";
							// System.out.println(sql);
							s.execute(sql);
						}
					} catch (Exception E) {
						E.printStackTrace();
					}
				}
				line = in.readLine();
			}
			// in.close();
			System.out.println("ProofPoint complete. ");

			// shutdown the various connections/file descriptors
			dis.close();
			dos.close();
			s.close();
			s2.close();
			db.closeConnection();
		} catch (Exception E) {
			E.printStackTrace();
		}
	}

}
